import numpy as np
from DataControl.getDataFromTushare.tudata import pullDataFromTushare
from sklearn.cluster import MiniBatchKMeans, KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.manifold import TSNE
from globleConfig.database import Database
import tushare as ts
from mplfinance.original_flavor import candlestick_ohlc
import matplotlib.dates as mpl_dates
from DataControl.orm.pgsql import pgsql
from matplotlib.dates import DateFormatter
import matplotlib.pyplot as plt
from sklearn.cluster import Birch
# from fastdtw import fastdtw
from scipy.spatial.distance import euclidean
from sklearn.cluster import AgglomerativeClustering
import pandas as pd
from DataControl.getDataFromTushare.tudata import pullDataFromTushare
from Index.MACDPoint import MACD
import datetime
from sklearn.cluster import Birch
from matplotlib.font_manager import FontProperties
from datetime import date

class clustering(object):
    """
    股票数据聚类类，用于对股票进行KMeans聚类分析。

    主要步骤包括：
    1. 从数据源获取所有股票的基本数据（如换手率和交易量）。
    2. 执行初步的KMeans聚类，依据换手率将股票划分为若干类别。
    3. 计算每个初步类别的平均交易量，并将初步类别根据平均交易量进行分组。
    4. 对每个初步类别的子集执行第二次KMeans聚类，根据交易量进一步细分股票类别。
    5. 将最终的聚类结果保存到CSV文件中。
    """
    def __init__(self):
        """
        初始化方法。设置默认参数并获取数据，开始初步的KMeans聚类。
        """
        self.pull = pullDataFromTushare()  # 实例化数据获取类
        self.allts_code = self.pull.getAllTs_code()  # 获取所有股票代码
        self.numofclass = 3  # 初始KMeans聚类类别数量
        self.numofseclass = 10  # 次级KMeans聚类类别数量
        self.data = dict()  # 存储所有股票数据的字典
        self.X_scaler = None  # 特征缩放器（待用）
        self.scalerofvol = None  # 交易量缩放器（待用）
        self.clusters = None  # 存储聚类结果
        self.features, self.turnover_rate, self.vol = self.dataprovider()  # 获取特征、换手率和交易量
        self.stock_code = list(self.data.keys())  # 股票代码列表
        self.clusteringt = None  # 初步聚类结果的DataFrame
        self.cluster_df = None  # 最终聚类结果的DataFrame
        self.kmeans1step()  # 执行初步KMeans聚类

    def kmeans1step(self):
        """
        执行初步的KMeans聚类步骤，根据换手率进行聚类，并根据交易量对每个类别进行进一步的KMeans聚类。
        """
        kmeansofturnover_rate = KMeans(n_clusters=self.numofclass)  # 初始化KMeans聚类器，用于换手率
        self.turnover_rate = np.array(self.turnover_rate).reshape(-1, 1)  # 将换手率数据转换为二维数组
        kmeansofturnover_rate.fit(self.turnover_rate)  # 对换手率数据进行聚类
        labelsoft = kmeansofturnover_rate.labels_  # 获取聚类标签
        self.clusteringt = self.label2df(labelsoft, self.stock_code)  # 将聚类标签转换为DataFrame

        # 计算每个类别的平均交易量
        label2meanoft = self._compute_mean_turnover_rate()  # 获取每个类别的平均交易量
        label2meanoft = dict(sorted(label2meanoft.items(), key=lambda item: item[1]))  # 按平均交易量排序

        # 将类别按平均交易量分组
        clusters = self._split_clusters(label2meanoft)  # 根据排序后的类别划分子集
        self._perform_secondary_clustering(clusters)  # 对每个子集进行第二次KMeans聚类

        self.cluster_df = self.clusteringt  # 更新最终聚类结果的DataFrame
        self._assign_final_labels()  # 将最终的聚类结果分配给DataFrame
        self.cluster_df.to_csv("cluster.csv", index=False)  # 将结果保存到CSV文件

    def _compute_mean_turnover_rate(self):
        """
        计算每个类别的平均交易量。
        """
        label2meanoft = {}  # 存储每个类别的平均交易量
        for i in range(self.numofclass):
            ts_codes = self.clusteringt.loc[self.clusteringt['cluster'] == i, 'ts_code']  # 获取属于该类别的股票代码
            mean_turnover_rate = np.mean([self.data[code]['vol'].sum() for code in ts_codes])  # 计算平均交易量
            label2meanoft[i] = mean_turnover_rate  # 存储到字典中
        return label2meanoft

    def _split_clusters(self, label2meanoft):
        """
        将类别按平均交易量分组，返回分组后的DataFrame。
        """
        tmpnumpy = np.array(list(label2meanoft.items()))  # 将字典转换为NumPy数组
        d1, d2, d3 = tmpnumpy[0:1, :], tmpnumpy[1:2, :], tmpnumpy[2:3, :]  # 划分子集
        df1, df2, df3 = self.splitnp2df(d1), self.splitnp2df(d2), self.splitnp2df(d3)  # 将子集转换为DataFrame
        return df1, df2, df3  # 返回子集的DataFrame

    def _perform_secondary_clustering(self, clusters):
        """
        对每个子集进行第二次KMeans聚类。
        """
        df1, df2, df3 = clusters  # 获取子集DataFrame
        secluster1, secluster2, secluster3 = self.kmeans2step(df1), self.kmeans2step(df2), self.kmeans2step(df3)  # 对每个子集执行KMeans聚类
        self.addmean(secluster1, 'cluster1')  # 为每个子集添加平均交易量和换手率
        self.addmean(secluster2, 'cluster2')
        self.addmean(secluster3, 'cluster3')

    def _assign_final_labels(self):
        """
        将最终的聚类标签分配给DataFrame。
        """
        zeros = np.zeros(self.cluster_df.shape[0])  # 初始化结果列
        self.cluster_df['res'] = zeros  # 添加结果列
        self.cluster_df.iloc[:, 2] = self.cluster_df.iloc[:, 2].astype(np.int64)  # 将聚类标签列转换为整数类型
        tmpdict = {}  # 存储最终聚类标签的字典
        for df in [self.secluster1, self.secluster2, self.secluster3]:  # 遍历每个子集DataFrame
            for i in range(df.shape[0]):
                tmpdict[df.iloc[i, 0]] = df.iloc[i, 1]  # 将子集的标签存储到字典中
        for i in range(self.cluster_df.shape[0]):
            self.cluster_df.iloc[i, 2] = int(tmpdict[self.cluster_df.iloc[i, 0]])  # 为每只股票分配最终的聚类标签

    def addmean(self, df, name):
        """
        为每个类别的DataFrame添加平均交易量和平均换手率，并将结果保存到CSV文件。
        """
        zero = np.zeros(df.shape[0])  # 初始化列
        df['volmean'] = zero  # 添加平均交易量列
        df['tmean'] = zero  # 添加平均换手率列
        for i in range(df.shape[0]):
            df.loc[i, 'volmean'] = self.data[df.loc[i, 'ts_code']].loc[:, 'vol'].mean()  # 计算并填充平均交易量
            df.loc[i, 'tmean'] = self.data[df.loc[i, 'ts_code']].loc[:, 'turnover_rate'].mean()  # 计算并填充平均换手率
        df.to_csv("clusters/" + name + ".csv", index=False)  # 保存结果到CSV文件

    def kmeans2step(self, df):
        """
        在每个子集上执行第二次KMeans聚类。
        """
        scaler = StandardScaler()  # 初始化特征缩放器
        stocks = df['ts_code'].values  # 获取股票代码
        vol = df['volmean'].values  # 获取交易量
        vol = vol.reshape(-1, 1)  # 将交易量转换为二维数组
        vol = scaler.fit_transform(vol)  # 对交易量进行标准化
        kmeansofvol = KMeans(n_clusters=self.numofseclass)  # 初始化第二次KMeans聚类器
        kmeansofvol.fit(vol)  # 对交易量数据进行聚类
        label = kmeansofvol.labels_  # 获取聚类标签
        res = self.label2df(label, stocks)  # 将标签转换为DataFrame
        return res

    def splitnp2df(self, d):
        """
        将NumPy数组转换为DataFrame。
        """
        res = pd.DataFrame()  # 初始化空DataFrame
        for i in range(d.shape[0]):
            cluster = d[i][0]  # 获取类别编号
            ts_codes = self.clusteringt.loc[self.clusteringt['cluster'] == cluster, 'ts_code']  # 获取属于该类别的股票代码
            for j in range(ts_codes.shape[0]):
                ts_code = ts_codes.iloc[j]  # 获取股票代码
                newrow = pd.Series([ts_code, self.data[ts_code]['vol'].mean()], ['ts_code', 'volmean'])  # 创建新行
                res = res.append(newrow, ignore_index=True)  # 将新行添加到DataFrame中
        return res

    def label2df(self, labels, stock_code):
        """
        将聚类标签转换为DataFrame。
        """
        clustered_stocks = {}  # 存储聚类结果的字典
        for stk_code, cluster_id in zip(stock_code, labels):  # 遍历每只股票和其对应的聚类标签
            if cluster_id not in clustered_stocks:
                clustered_stocks[cluster_id] = []  # 初始化类别列表
            clustered_stocks[cluster_id].append(stk_code)  # 将股票代码添加到类别列表中

        cluster_res = []  # 存储最终结果的列表
        for cluster_id, stk in clustered_stocks.items():  # 遍历每个类别
            for stock in stk:  # 遍历每只股票
                cluster_res.append({'ts_code': stock, 'cluster': cluster_id})  # 将股票代码和类别编号添加到结果列表中
        cluster_df = pd.DataFrame(cluster_res)  # 将结果列表转换为DataFrame
        return cluster_df

    def dataprovider(self):
        """
        从数据源获取股票数据，并计算每只股票的特征（换手率和交易量的平均值）。
        """
        for i in range(self.allts_code.shape[0]):
            ts_code = self.allts_code[i]  # 获取股票代码
            print(ts_code)  # 打印当前处理的股票代码
            num = self.pull.get_count_daily(ts_code)  # 获取该股票的日数据条数
            tmpdata = self.pull.get_lastk_daily(ts_code=ts_code, k=min(num, 100))  # 获取最近100条或所有数据
            tmpdata = tmpdata[['vol', 'turnover_rate']]  # 提取交易量和换手率数据
            if tmpdata.shape[0] == 0:  # 如果数据为空，则跳过
                continue


if __name__ == '__main__':
    pd.set_option('display.max_columns', None)

    # 设置显示的最大行为不限制
    pd.set_option('display.max_rows', None)

    # 设置最大宽度，避免列宽带来的截断，None表示自动检测
    pd.set_option('display.width', None)

    # 设置每列的最大宽度，避免列内容被截断，None表示不限制
    pd.set_option('display.max_colwidth', None)
    cluster = clustering()
    #cluster.clustered_pic()
















